from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import argparse
import backtrader as bt
import backtrader.feeds as btfeeds
from DB.FuturesDB import getFutureMinPrice, getFutureDailyPrice
import pandas


# 获取数据的频率跟策略执行的一致
def build_backtest_dataset(start_date:str, end_date:str):
    """构造配套 backtesting 用的数据集"""
    # Backtest 的数据源为一个包含 4 列 ['Open', 'High', 'Low', 'Close'] 的 dataframe
    df = getFutureMinPrice(start_date, end_date, 'T9999_Min', format=1)
    df = df[['open', 'high', 'low', 'close', 'volume']]
    # df.columns = ['Open', 'High', 'Low', 'Close']
    df = df.sort_index()

    return df


def runstrat():
    args = parse_args()

    # Create a cerebro entity
    cerebro = bt.Cerebro(stdstats=False)

    # Add a strategy
    cerebro.addstrategy(bt.Strategy)

    # # Get a pandas dataframe
    # datapath = ('../../datas/2006-day-001.txt')
    #
    # # Simulate the header row isn't there if noheaders requested
    # skiprows = 1 if args.noheaders else 0
    # header = None if args.noheaders else 0
    #
    # dataframe = pandas.read_csv(datapath,
    #                             skiprows=skiprows,
    #                             header=header,
    #                             parse_dates=True,
    #                             index_col=0)
    train_start = '2019-01-07'
    train_end = '2019-12-31'
    dataframe = build_backtest_dataset(train_start, train_end)

    if not args.noprint:
        print('--------------------------------------------------')
        print(dataframe)
        print('--------------------------------------------------')

    # Pass it to the backtrader datafeed and add it to the cerebro
    # data = bt.feeds.PandasData(dataname=dataframe, open=0, high=1, low=2, close=3, volume=4)
    # data = bt.feeds.PandasData(dataname=dataframe, datetime=0, open=1, high=2, low=3, close=4, volume=5)
    data = bt.feeds.PandasData(dataname=dataframe)
    # 使用PandasDirectData代替PandasData加快数据加载，加载速度能提升大约1倍。
    # data = bt.feeds.PandasDirectData(dataname=dataframe, open=0, high=1, low=2, close=3, volume=4)
    cerebro.adddata(data)

    # Run over everything
    cerebro.run()

    # Plot the result
    cerebro.plot(style='candlestick')


def parse_args():
    parser = argparse.ArgumentParser(
        description='Pandas test script')

    parser.add_argument('--noheaders', action='store_true', default=False,
                        required=False,
                        help='Do not use header rows')

    parser.add_argument('--noprint', action='store_true', default=False,
                        help='Print the dataframe')

    return parser.parse_args()


if __name__ == '__main__':
    runstrat()